{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Table of Contents\n",
    "* [Intro](#Intro)\n",
    "\t* [Generative Adversarial Networks (GANs)](#Generative-Adversarial-Networks-%28GANs%29)\n",
    "* [Gaussian Distribution Approximation (Keras)](#Gaussian-Distribution-Approximation-%28Keras%29)\n",
    "* [MNIST GAN (Keras)](#MNIST-GAN-%28Keras%29)\n",
    "\t* [Generator Model](#Generator-Model)\n",
    "\t* [Discriminator Model](#Discriminator-Model)\n",
    "\t* [GAN Model](#GAN-Model)\n",
    "\t* [Train](#Train)\n",
    "* [Wasserstein GAN](#Wasserstein-GAN)\n",
    "\t* [Generator Model](#Generator-Model)\n",
    "\t* [Discriminator Model](#Discriminator-Model)\n",
    "\t* [GAN Model](#GAN-Model)\n",
    "\t* [Train](#Train)\n",
    "\t* [Plot Losses](#Plot-Losses)\n",
    "* [Generate gif](#Generate-gif)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Intro"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Exploratory notebook related to Generative Adversarial Networks (GANs). Includes toy examples implementation and testing of related techniques or subjects."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generative Adversarial Networks (GANs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Architecture that learns by posing two networks in competition with each others. Goal is to learn parameters in order to produce a distribution close to our dataset distribution (true distribution).\n",
    "* Discriminator: detect if an image belongs to a target dataset or not (generated by the generator) \n",
    "* Generator: generate new examples that looks like the training/target data (fool the discriminator)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2017-10-24T15:27:18.692240",
     "start_time": "2017-10-24T15:27:18.009201Z"
    }
   },
   "outputs": [],
   "source": [
    "import time\n",
    "from PIL import Image\n",
    "import numpy as np\n",
    "import pdb\n",
    "import os\n",
    "import sys\n",
    "import seaborn as sns\n",
    "\n",
    "import matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib import animation\n",
    "\n",
    "from keras.models import Sequential\n",
    "from keras.layers.core import Reshape, Dense, Dropout, Activation, Flatten\n",
    "from keras.layers.convolutional import Convolution2D, MaxPooling2D, ZeroPadding2D, UpSampling2D\n",
    "from keras.layers.advanced_activations import LeakyReLU\n",
    "from keras import backend as K\n",
    "from keras import optimizers\n",
    "from keras.layers.normalization import BatchNormalization\n",
    "from keras.datasets import mnist\n",
    "\n",
    "from tqdm import tqdm_notebook as tqdm\n",
    "\n",
    "%matplotlib notebook\n",
    "sns.set_style(\"dark\")\n",
    "\n",
    "\n",
    "sys.path.append(os.path.join(os.getcwd(), os.pardir))\n",
    "from utils.plot_utils import plot_sample_imgs\n",
    "from utils.generative_utils import NoiseDistribution, set_trainable\n",
    "\n",
    "RES_DIR = os.path.join(*[os.pardir]*2, 'data', 'deep_learning')\n",
    "\n",
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "# Gaussian Distribution Approximation (Keras)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "Example adapted from [Aylien blog](http://blog.aylien.com/introduction-generative-adversarial-networks-code-tensorflow/).\n",
    "\n",
    "Check also [here](https://medium.com/towards-data-science/gan-by-example-using-keras-on-tensorflow-backend-1a6d515a60d0) for Keras code"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2017-07-17T14:38:35.436225",
     "start_time": "2017-07-17T14:38:35.354220Z"
    },
    "collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# target 1D gaussian distribution class\n",
    "class GaussianDistribution:\n",
    "    def __init__(self, mu=4, sigma=0.5):\n",
    "        self.mu = mu\n",
    "        self.sigma = sigma\n",
    "    \n",
    "    def sample(self, N):\n",
    "        samples = np.random.normal(self.mu, self.sigma, N)\n",
    "        samples.sort()\n",
    "        return samples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2017-07-17T14:38:36.371278",
     "start_time": "2017-07-17T14:38:36.362277Z"
    },
    "collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# generator input noise distribution class\n",
    "class GeneratorNoiseDistribution:\n",
    "    def __init__(self, vrange):\n",
    "        self.vrange = vrange\n",
    "        \n",
    "    def sample(self, N):\n",
    "        return np.linspace(-self.vrange, self.vrange, N) + \\\n",
    "                    np.random.random(N) * 0.01"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2017-07-17T14:39:01.626723",
     "start_time": "2017-07-17T14:39:01.620722Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def generator(input_dim, hidden_size):\n",
    "    g = Sequential()\n",
    "    g.add(Dense(hidden_size, input_dim=input_dim, activation=K.softplus))\n",
    "    g.add(Dense(input_dim))\n",
    "    return g"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2017-07-17T14:39:01.641723",
     "start_time": "2017-07-17T14:39:01.627723Z"
    },
    "collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def discriminator(input_dim, hidden_size):\n",
    "    d = Sequential()\n",
    "    d.add(Dense(hidden_size*2, input_dim=input_dim, activation=K.tanh))\n",
    "    d.add(Dense(hidden_size*2, activation=K.tanh))\n",
    "    d.add(Dense(hidden_size*2, activation=K.tanh))\n",
    "    d.add(Dense(1, activation=K.sigmoid))\n",
    "    return d"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2017-07-17T15:18:45.655081",
     "start_time": "2017-07-17T15:18:45.646081Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# init distributions\n",
    "gaussian_d = GaussianDistribution()\n",
    "generator_d = GeneratorNoiseDistribution(8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2017-07-17T15:18:45.876094",
     "start_time": "2017-07-17T15:18:45.791089Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# init GAN components\n",
    "d = discriminator(1, 128)\n",
    "g = generator(1, 128)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2017-07-17T15:18:47.018159",
     "start_time": "2017-07-17T15:18:46.961156Z"
    },
    "collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# discriminator model\n",
    "optimizer = optimizers.RMSprop(lr=0.0008, clipvalue=1.0, decay=6e-8)\n",
    "discriminator_model = d\n",
    "discriminator_model.compile(loss='binary_crossentropy', optimizer=optimizer, metrics=['accuracy'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2017-07-17T15:18:49.135280",
     "start_time": "2017-07-17T15:18:48.930268Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# adversarial model\n",
    "optimizer = optimizers.RMSprop(lr=0.0004, clipvalue=1.0, decay=3e-8)\n",
    "adversarial_model = Sequential()\n",
    "adversarial_model.add(g)\n",
    "adversarial_model.add(d)\n",
    "adversarial_model.compile(loss='binary_crossentropy', optimizer=optimizer, metrics=['accuracy'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2017-07-17T15:18:53.961556",
     "start_time": "2017-07-17T15:18:53.826549Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "batch_size = 64\n",
    "\n",
    "fig, ax = plt.subplots(dpi=100, figsize=(5, 4))\n",
    "true_dist = np.reshape(gaussian_d.sample(1000), (1000, 1))\n",
    "plt.show()\n",
    "\n",
    "def animate(step):\n",
    "#for step in range(100):\n",
    "    \n",
    "    # generate data\n",
    "    # first we sample from the true distribution, then we generate some\n",
    "    # \"fake\" data by feeding noise to the generator\n",
    "    true_sample = np.reshape(gaussian_d.sample(batch_size), (batch_size, 1))\n",
    "    noise = generator_d.sample(batch_size)\n",
    "    fake_sample = g.predict(noise)\n",
    "    #pdb.set_trace()\n",
    "    \n",
    "    # train discriminator\n",
    "    # feed true and fake samples with respective labels (1, 0) to the discriminator\n",
    "    x = np.reshape(np.concatenate((true_sample, fake_sample)), (batch_size*2, 1))\n",
    "    y = np.ones([batch_size*2, 1])\n",
    "    y[batch_size:, :] = 0\n",
    "    d_loss = discriminator_model.train_on_batch(x, y)\n",
    "    \n",
    "    # train GAN\n",
    "    # feed noise to the model and expect true (1) response from discriminator,\n",
    "    # which is in turn fed with data generated by the generator\n",
    "    noise = np.reshape(generator_d.sample(batch_size), (batch_size, 1))\n",
    "    y = np.ones([batch_size, 1])\n",
    "    a_loss = adversarial_model.train_on_batch(noise, y)\n",
    "    \n",
    "    log_mesg = \"%d: [D loss: %f, acc: %f]\" % (step, d_loss[0], d_loss[1])\n",
    "    log_mesg = \"%s  [A loss: %f, acc: %f]\" % (log_mesg, a_loss[0], a_loss[1])\n",
    "    \n",
    "    # plot\n",
    "    fig.clf()\n",
    "    fake = sns.distplot(fake_sample)\n",
    "    fake.set_xlim([0,8])\n",
    "    fake.set_ylim([0,3])\n",
    "    sns.distplot(true_dist)\n",
    "    sns.plt.text(3, 2, \"Epoch {}, a_loss {:.3f}\".format(step, a_loss[0]))\n",
    "    \n",
    "anim = animation.FuncAnimation(fig, animate, 200, repeat=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2017-05-28T13:30:43.893038",
     "start_time": "2017-05-28T13:30:43.885037Z"
    },
    "collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "noise = generator_d.sample(batch_size)\n",
    "fake_sample = g.predict(noise)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2017-05-28T13:30:45.265116",
     "start_time": "2017-05-28T13:30:45.137109Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "sns.distplot(fake_sample)\n",
    "sns.plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "# MNIST GAN (Keras)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Example adapted from [MNIST Generative Adversarial Model in Keras](http://www.kdnuggets.com/2016/07/mnist-generative-adversarial-model-keras.html)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2017-10-24T08:07:35.586512",
     "start_time": "2017-10-24T08:07:34.371443Z"
    },
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "noise_d = NoiseDistribution()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2017-10-20T09:18:05.565305",
     "start_time": "2017-10-20T09:18:03.881209Z"
    },
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "input_dim = 100\n",
    "img_shape = (28,28,1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generator Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2017-10-16T11:22:55.799545",
     "start_time": "2017-10-16T11:22:55.777544Z"
    },
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# model takes real values vector of size input_dim and via upsampling,\n",
    "# reshaping, and various convolutional filters generates a 28x28 b/w image\n",
    "def generator_model(input_dim, n_channels=128, init_side=7):\n",
    "    m = Sequential()\n",
    "    m.add(Dense(init_side*init_side*n_channels, input_dim=input_dim, activation=LeakyReLU()))\n",
    "    m.add(BatchNormalization(mode=2))\n",
    "    m.add(Reshape((init_side, init_side, n_channels)))\n",
    "    \n",
    "    m.add(UpSampling2D())\n",
    "    m.add(Convolution2D(n_channels//2, 3, 3, border_mode='same', activation=LeakyReLU()))\n",
    "    m.add(BatchNormalization(mode=2))\n",
    "    m.add(UpSampling2D())\n",
    "    m.add(Convolution2D(n_channels//4, 3, 3, border_mode='same', activation=LeakyReLU()))\n",
    "    m.add(BatchNormalization(mode=2))\n",
    "    #?? Tanh \n",
    "    m.add(Convolution2D(1, 1, 1, border_mode='same', activation='sigmoid'))\n",
    "    return m"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2017-10-16T11:22:56.928610",
     "start_time": "2017-10-16T11:22:56.546588Z"
    }
   },
   "outputs": [],
   "source": [
    "g = generator_model(input_dim=input_dim, n_channels=512)\n",
    "g.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2017-10-16T11:22:24.556758",
     "start_time": "2017-10-16T11:22:22.171622Z"
    },
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# plot random generated image\n",
    "plt.imshow(g.predict(noise_d.sample((1, input_dim)))[0]\n",
    "           .reshape(28, 28))\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Discriminator Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2017-10-16T11:23:01.252857",
     "start_time": "2017-10-16T11:23:01.239856Z"
    }
   },
   "outputs": [],
   "source": [
    "# model takes image and after convolution and flattening\n",
    "# outputs a probability value\n",
    "def discriminator_model(input_shape, init_filters=64):\n",
    "    m = Sequential()\n",
    "    m.add(Convolution2D(init_filters, 5, 5, subsample=(2, 2), input_shape=input_shape, border_mode='same',\n",
    "                       activation=LeakyReLU(0.2)))\n",
    "    #?? maxpooling and dropout? MaxPool2D(pool_size=2)\n",
    "    m.add(Convolution2D(init_filters*2, 5, 5, subsample=(2, 2), border_mode='same',\n",
    "                       activation=LeakyReLU(0.2)))\n",
    "    #m.add(Convolution2D(init_filters*4, 3, 5, border_mode='same', \n",
    "    #                    activation=LeakyReLU(0.2)))\n",
    "    m.add(Flatten())\n",
    "    m.add(Dense(256, activation=LeakyReLU()))\n",
    "    m.add(Dense(1, activation='sigmoid'))\n",
    "    return m"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2017-10-16T11:23:22.256058",
     "start_time": "2017-10-16T11:23:22.129051Z"
    }
   },
   "outputs": [],
   "source": [
    "d = discriminator_model(input_shape=(28,28,1), init_filters=256)\n",
    "d.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2017-10-16T11:24:01.201286",
     "start_time": "2017-10-16T11:24:00.957272Z"
    }
   },
   "outputs": [],
   "source": [
    "# print prediction for random image\n",
    "d.predict(g.predict(noise_d.sample((1, input_dim))))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## GAN Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2017-10-16T11:30:42.280226",
     "start_time": "2017-10-16T11:30:41.702193Z"
    }
   },
   "outputs": [],
   "source": [
    "# init GAN components\n",
    "g = generator_model(input_dim)\n",
    "d = discriminator_model(img_shape)\n",
    "\n",
    "# compile generator\n",
    "#g_optimizer = optimizers.Adam(lr=0.0001)\n",
    "#g.compile(loss='binary_crossentropy', optimizer=g_optimizer)\n",
    "\n",
    "# compile discriminator\n",
    "d_optimizer = optimizers.Adam(lr=0.001)\n",
    "d.compile(loss='binary_crossentropy', optimizer=d_optimizer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2017-10-16T11:30:42.804256",
     "start_time": "2017-10-16T11:30:42.698250Z"
    },
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# build adversarial model\n",
    "gan = Sequential()\n",
    "gan.add(g)\n",
    "gan.add(d)\n",
    "gan_optimizer = optimizers.Adam(lr=0.0001)\n",
    "gan.compile(loss='binary_crossentropy', optimizer=gan_optimizer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2017-10-16T11:24:19.013305",
     "start_time": "2017-10-16T11:24:19.005304Z"
    }
   },
   "outputs": [],
   "source": [
    "gan.summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "generator_fun = lambda num_samples : generator.predict(noise_d.sample((num_samples, input_dim)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2017-10-16T11:30:49.736653",
     "start_time": "2017-10-16T11:30:48.254568Z"
    }
   },
   "outputs": [],
   "source": [
    "# load mnist data using Keras\n",
    "(X_train, y_train), (X_test, y_test) = mnist.load_data()\n",
    "# reshape and normalize train data\n",
    "X_train = np.expand_dims(X_train, axis=-1)\n",
    "X_train = X_train.astype('float32')/255\n",
    "print(X_train.shape)\n",
    "print(y_train.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2017-10-16T11:30:49.762654",
     "start_time": "2017-10-16T11:30:49.739653Z"
    },
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def train_discriminator(d, g, noise_d, input_dim, X_train, batch_size, epoch):\n",
    "    # generate data\n",
    "    # first we sample from the true distribution (mnist dataset), then we generate some\n",
    "    # \"fake\" images by feeding noise to the generator\n",
    "    \n",
    "    # generate random indexes for train data\n",
    "    batch_idxs = np.random.randint(0, len(X_train), batch_size)\n",
    "    # collect images corresponsing to previously generated index, and add a dimension\n",
    "    true_sample = X_train[batch_idxs,:,:,:]\n",
    "    # generate fake sample\n",
    "    fake_sample = g.predict(noise_d.sample((batch_size, input_dim)))\n",
    "    \n",
    "    # prepare train batch data\n",
    "    # concatenativ true and fake samples and adjusting labels accordingly\n",
    "    x = np.concatenate((true_sample, fake_sample))\n",
    "    y = np.ones([batch_size*2, 1])\n",
    "    y[batch_size:,:] = 0\n",
    "    \n",
    "    # train discriminator\n",
    "    # feed true and fake samples with respective labels (1, 0) to the discriminator\n",
    "    set_trainable(d, True, None, None)\n",
    "    d_loss = d.train_on_batch(x, y)\n",
    "    #print(\"Epoch {}: [D loss: {}]\".format(epoch, d_loss))\n",
    "    return d_loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2017-10-16T11:30:51.225738",
     "start_time": "2017-10-16T11:30:51.197736Z"
    }
   },
   "outputs": [],
   "source": [
    "def train(d, g, gan, noise_d, input_dim, X_train, batch_size=32, n_epochs=1, add_epoch=0):\n",
    "    losses = {'g':[], 'd':[]}\n",
    "    for epoch in range(n_epochs):\n",
    "        # train discriminator\n",
    "        d_loss = train_discriminator(d, g, noise_d, input_dim, X_train, batch_size, epoch)\n",
    "        losses['d'].append(d_loss)\n",
    "        set_trainable(d, False, None, None)\n",
    "\n",
    "        # train GAN\n",
    "        # feed noise to the model and expect true (1) response from discriminator,\n",
    "        # which is in turn fed with data generated by the generator\n",
    "        noise = noise_d.sample((batch_size, input_dim))\n",
    "        y = np.ones([batch_size, 1])\n",
    "\n",
    "        g_loss = gan.train_on_batch(noise, y)\n",
    "        losses['g'].append(g_loss)\n",
    "\n",
    "        #print(\"Epoch {}: [G loss: {}]\".format(epoch, g_loss))\n",
    "\n",
    "        if (epoch%10)==0:\n",
    "            plot_sample_imgs(generator_fun, img_shape[:2], savepath=os.path.join('data', 'mnist_gan', 'mnist_gen{}.jpg'.format(epoch+add_epoch)))\n",
    "    return losses"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2017-10-16T11:30:55.751997",
     "start_time": "2017-10-16T11:30:54.173906Z"
    }
   },
   "outputs": [],
   "source": [
    "# pretrain discriminator\n",
    "batch_size = 128\n",
    "n_epochs = 1\n",
    "for epoch in range(n_epochs):\n",
    "    train_discriminator(d, g, noise_d, input_dim, X_train, batch_size, epoch)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2017-10-16T11:59:09.424869",
     "start_time": "2017-10-16T11:35:23.673321Z"
    }
   },
   "outputs": [],
   "source": [
    "#plt.ion()\n",
    "plt.ioff()\n",
    "\n",
    "K.set_value(d.optimizer.lr, 1e-3)\n",
    "K.set_value(gan.optimizer.lr, 1e-3)\n",
    "\n",
    "losses = train(d, g, gan, noise_d, input_dim, X_train, \n",
    "               batch_size=256, n_epochs=1000, add_epoch=120)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2017-07-26T19:16:17.315164",
     "start_time": "2017-07-26T19:16:17.199158Z"
    }
   },
   "outputs": [],
   "source": [
    "# plot random generated image\n",
    "plt.imshow(g.predict(np.random.randn(input_dim).reshape(1, -1))\n",
    "           .reshape(28, 28))\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2017-07-25T16:16:20.865150",
     "start_time": "2017-07-25T16:16:20.318119Z"
    }
   },
   "outputs": [],
   "source": [
    "plt.imshow(true_sample[2].reshape(28, 28))\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2017-07-26T19:18:08.697535",
     "start_time": "2017-07-26T19:18:03.669247Z"
    }
   },
   "outputs": [],
   "source": [
    "gan.test_on_batch(noise, y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2017-07-25T16:34:49.995589",
     "start_time": "2017-07-25T16:34:46.581393Z"
    }
   },
   "outputs": [],
   "source": [
    "gan.train_on_batch(noise, y)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Wasserstein GAN"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[Source](https://myurasov.github.io/2017/09/24/wasserstein-gan-keras.html)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2017-10-24T08:25:46.282896",
     "start_time": "2017-10-24T08:25:44.881816Z"
    }
   },
   "outputs": [],
   "source": [
    "from keras.layers import *\n",
    "from keras.models import *\n",
    "from keras.optimizers import *\n",
    "from keras.initializers import *\n",
    "from keras.callbacks import *\n",
    "from keras.utils.generic_utils import Progbar"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2017-10-24T08:25:47.746980",
     "start_time": "2017-10-24T08:25:46.283896Z"
    },
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "noise_d = NoiseDistribution()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2017-10-24T08:25:49.136060",
     "start_time": "2017-10-24T08:25:47.748980Z"
    },
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "input_dim = 100\n",
    "img_shape = (28,28,1)\n",
    "num_classes = 10"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generator Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2017-10-24T08:25:50.645146",
     "start_time": "2017-10-24T08:25:49.138060Z"
    },
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# utility for the standard deconvolution block used in the generator\n",
    "def generator_deconv_block(filters, block_input, kernel_size=(3, 3), strides=(1, 1)):\n",
    "    block = UpSampling2D()(block_input)\n",
    "    block = Convolution2D(filters, (3, 3), strides=strides, padding='same')(block)\n",
    "    block = LeakyReLU()(block)\n",
    "    block = BatchNormalization()(block)\n",
    "    return block"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2017-10-24T08:25:52.073228",
     "start_time": "2017-10-24T08:25:50.648146Z"
    },
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# different from basic DCGAN, this WGAN model\n",
    "# takes as input both the prior sample (noise) and the image class\n",
    "def generator_model(input_dim, voc_size, init_filters=128, init_side=7, num_deconv_blocks=2):\n",
    "    # Input combination part\n",
    "    input_class = Input(shape=(1, ), dtype='int32')\n",
    "    e = Embedding(voc_size, input_dim)(input_class)\n",
    "    embedded_class = Flatten()(e)\n",
    "    # noise\n",
    "    noise = Input(shape=(input_dim, ))\n",
    "    # hadamard product\n",
    "    h = multiply([noise, embedded_class])\n",
    "    \n",
    "    # CNN part\n",
    "    x = Dense(1024)(h)\n",
    "    x = LeakyReLU()(x)\n",
    "    \n",
    "    x = Dense(init_side*init_side*init_filters)(x)\n",
    "    x = LeakyReLU()(x)\n",
    "    x = BatchNormalization()(x)\n",
    "    x = Reshape((init_side, init_side, init_filters))(x)\n",
    "\n",
    "    for i in range(num_deconv_blocks):\n",
    "        x = generator_deconv_block(init_filters//(2**(i+1)), block_input=x, kernel_size=(5, 5))\n",
    "\n",
    "    x = Convolution2D(1, (2, 2), padding='same', activation='tanh')(x)\n",
    "    \n",
    "    return Model(inputs=[noise, input_class], outputs=x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2017-10-24T08:25:54.034340",
     "start_time": "2017-10-24T08:25:52.075228Z"
    }
   },
   "outputs": [],
   "source": [
    "# instantiate generate model\n",
    "gen = generator_model(input_dim=input_dim, voc_size=10, init_filters=128)\n",
    "gen.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2017-10-24T11:32:51.132622",
     "start_time": "2017-10-24T11:32:49.700540Z"
    }
   },
   "outputs": [],
   "source": [
    "# plot random generated image\n",
    "plt.imshow(gen.predict([noise_d.sample((1, input_dim)), np.array([7])])[0]\n",
    "           .reshape(28, 28))\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Discriminator Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2017-10-24T08:25:59.042626",
     "start_time": "2017-10-24T08:25:57.649546Z"
    },
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# utility for the standard convolution block used in the discriminator\n",
    "def discriminator_conv_block(filters, block_input, kernel_size=(3, 3), strides=(1, 1), pool_size=None):\n",
    "    block = Convolution2D(filters, kernel_size, strides=strides, padding='same')(block_input)\n",
    "    block = LeakyReLU()(block)\n",
    "    block = BatchNormalization()(block)\n",
    "    \n",
    "    # if given, add max pooling\n",
    "    if pool_size:\n",
    "        block = MaxPool2D(pool_size=pool_size)(block)\n",
    "        \n",
    "    return block"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2017-10-24T08:26:00.414705",
     "start_time": "2017-10-24T08:25:59.044626Z"
    },
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# different from basic DCGAN, this WGAN discriminator model\n",
    "# takes an image as input, and output both a prediction about image autheticity\n",
    "# as well as one for the image class\n",
    "def discriminator_model(input_shape, num_classes, init_filters=32, num_conv_blocks=3):\n",
    "    input_image = Input(shape=input_shape)\n",
    "    \n",
    "    x = input_image\n",
    "    for i in range(num_conv_blocks):\n",
    "        x = discriminator_conv_block(init_filters*(2**i), block_input=x, pool_size=None)\n",
    "\n",
    "    features = Flatten()(x)\n",
    "    \n",
    "    out_autheticity = Dense(1, activation='linear')(features)\n",
    "    out_class = Dense(num_classes, activation='softmax')(features)\n",
    "    \n",
    "    return Model(inputs=[input_image], outputs=[out_autheticity, out_class])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2017-10-24T08:15:08.309406",
     "start_time": "2017-10-24T08:15:06.708315Z"
    },
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# instantiate discriminator model\n",
    "dis = discriminator_model(input_shape=img_shape, num_classes=10, init_filters=32)\n",
    "dis.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2017-10-24T08:15:13.210687",
     "start_time": "2017-10-24T08:15:11.702600Z"
    }
   },
   "outputs": [],
   "source": [
    "# print prediction for random image\n",
    "dis.predict(gen.predict([noise_d.sample((1, input_dim)), np.array([3])]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## GAN Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2017-10-24T08:26:08.486166",
     "start_time": "2017-10-24T08:26:07.082086Z"
    },
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# loss function for discriminator\n",
    "def d_loss(y_true, y_pred):\n",
    "    return K.mean(y_true * y_pred)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2017-10-24T08:26:10.699293",
     "start_time": "2017-10-24T08:26:08.663176Z"
    }
   },
   "outputs": [],
   "source": [
    "# init GAN components\n",
    "gen = generator_model(input_dim=input_dim, voc_size=num_classes, init_filters=128)\n",
    "dis = discriminator_model(input_shape=img_shape, num_classes=num_classes, init_filters=32)\n",
    "\n",
    "# compile discriminator\n",
    "dis.compile(loss=[d_loss, 'sparse_categorical_crossentropy'], \n",
    "            optimizer=RMSprop(lr=1e-4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2017-10-24T08:26:12.644404",
     "start_time": "2017-10-24T08:26:10.701293Z"
    }
   },
   "outputs": [],
   "source": [
    "# Build adversarial model\n",
    "\n",
    "noise = Input(shape=(input_dim, ))\n",
    "input_class = Input(shape=(1, ), dtype='int32')\n",
    "out_autheticity, out_class = dis(gen(inputs=[noise, input_class]))\n",
    "\n",
    "\n",
    "gan = Model(inputs=[noise, input_class], outputs=[out_autheticity, out_class])\n",
    "gan.compile(loss=[d_loss, 'sparse_categorical_crossentropy'], \n",
    "            optimizer=RMSprop(lr=1e-4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2017-10-22T09:23:52.912952",
     "start_time": "2017-10-22T09:23:51.370863Z"
    }
   },
   "outputs": [],
   "source": [
    "gan.summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2017-10-24T08:26:16.722637",
     "start_time": "2017-10-24T08:26:15.320557Z"
    },
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# given that generator uses tanh activation function, \n",
    "# we need to process its output to make it a valid image\n",
    "#deprocess = lambda img : np.transpose((img/2+0.5).clip(0,1), (1,2,0))\n",
    "deprocess = lambda img : (img/2+0.5).clip(0,1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2017-10-24T08:26:18.107717",
     "start_time": "2017-10-24T08:26:16.724638Z"
    },
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "generator_fun = lambda num_samples : deprocess(gen.predict([noise_d.sample((num_samples, input_dim)),\n",
    "                               np.random.randint(0, num_classes, num_samples)]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2017-10-24T08:26:20.032827",
     "start_time": "2017-10-24T08:26:18.109717Z"
    }
   },
   "outputs": [],
   "source": [
    "# load mnist data using Keras\n",
    "(X_train, y_train), (X_test, y_test) = mnist.load_data()\n",
    "# normalize to -1..1 range and reshape\n",
    "X_train = (X_train.astype(np.float32) - 127.5) / 127.5\n",
    "X_train = np.expand_dims(X_train, axis=-1)\n",
    "print(X_train.shape)\n",
    "print(y_train.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2017-10-24T08:26:23.929050",
     "start_time": "2017-10-24T08:26:22.466966Z"
    }
   },
   "outputs": [],
   "source": [
    "def train_discriminator(dis, gen, noise_d, input_dim, num_classes, X_train, Y_train, batch_size, epoch):\n",
    "    # clip weights\n",
    "    for l in dis.layers:\n",
    "        weights = l.get_weights()\n",
    "        weights = [np.clip(w, -0.01, 0.01) for w in weights]\n",
    "        l.set_weights(weights)\n",
    "    \n",
    "    # generate data\n",
    "    # first we sample from the true distribution (mnist dataset), then we generate some\n",
    "    # \"fake\" images by feeding noise to the generator\n",
    "    \n",
    "    # generate random indexes for train data\n",
    "    batch_idxs = np.random.randint(0, len(X_train), batch_size)\n",
    "    # collect images corresponsing to previously generated index, and add a dimension\n",
    "    true_sample = X_train[batch_idxs]\n",
    "    true_sample_classes = y_train[batch_idxs]\n",
    "    \n",
    "    # train on true samples\n",
    "    dis_true_loss = dis.train_on_batch(true_sample, \n",
    "                                       [-np.ones(batch_size), true_sample_classes])\n",
    "    \n",
    "    # generate fake sample\n",
    "    noise = noise_d.sample((batch_size, input_dim))\n",
    "    generated_classes = np.random.randint(0, num_classes, batch_size)\n",
    "    fake_sample = gen.predict([noise, generated_classes.reshape(-1, 1)])\n",
    "    \n",
    "    # train on fake samples\n",
    "    dis_fake_loss = dis.train_on_batch(fake_sample, \n",
    "                                   [np.ones(batch_size), generated_classes])\n",
    "    #print(\"Epoch {}: [D loss: {}]\".format(epoch, d_loss))\n",
    "    return dis_true_loss, dis_fake_loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2017-10-24T08:26:41.104032",
     "start_time": "2017-10-24T08:26:39.590945Z"
    },
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def train(dis, gen, gan, noise_d, input_dim, num_classes, X_train, Y_train,\n",
    "          batch_size=32, n_epochs=1, add_epochs=0):\n",
    "    losses = {'gan':[], 'dis_fake_loss':[], 'dis_true_loss':[]}\n",
    "    for epoch in tqdm(range(n_epochs), desc='Training GAN'):\n",
    "        if (epoch+add_epochs % 1000) < 15 or epoch+add_epochs % 500 == 0: # 25 times in 1000, every 500th\n",
    "            d_iters = 40\n",
    "        else:\n",
    "            d_iters = 5#D_ITERS\n",
    "        \n",
    "        # train discriminator\n",
    "        set_trainable(dis, True, None, None)\n",
    "        for d_epoch in range(d_iters):\n",
    "            dis_true_loss, dis_fake_loss = train_discriminator(dis, gen, noise_d, input_dim, num_classes, \n",
    "                                         X_train, Y_train, batch_size, epoch)\n",
    "            losses['dis_fake_loss'].append(dis_fake_loss)\n",
    "            losses['dis_true_loss'].append(dis_true_loss)\n",
    "        \n",
    "        set_trainable(dis, False, None, None)\n",
    "\n",
    "        # train GAN\n",
    "        # feed noise to the model and expect true (1) response from discriminator,\n",
    "        # which is in turn fed with data generated by the generator\n",
    "        noise = noise_d.sample((batch_size, input_dim))\n",
    "        generated_classes = np.random.randint(0, num_classes, batch_size)\n",
    "\n",
    "        gan_loss = gan.train_on_batch(\n",
    "            [noise, generated_classes.reshape((-1, 1))], \n",
    "            [-np.ones(batch_size), generated_classes])\n",
    "        losses['gan'].append(gan_loss)\n",
    "\n",
    "        #print(\"Epoch {}: [G loss: {}]\".format(epoch, g_loss))\n",
    "\n",
    "        if epoch%10 == 0:\n",
    "            plot_sample_imgs(generator_fun, img_shape[:2], \n",
    "                             savepath=os.path.join(RES_DIR, 'data', 'mnist_wgan', 'mnist_gen{}.jpg'.format(epoch+add_epochs)))\n",
    "    return losses"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2017-10-24T08:26:43.732182",
     "start_time": "2017-10-24T08:26:42.358104Z"
    },
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "add_epochs = 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2017-10-24T13:36:18.704310",
     "start_time": "2017-10-24T12:49:01.682042Z"
    }
   },
   "outputs": [],
   "source": [
    "#plt.ion()\n",
    "plt.ioff()\n",
    "\n",
    "n_epochs = 500\n",
    "losses = train(dis, gen, gan, noise_d, input_dim, num_classes, \n",
    "               X_train, y_train, \n",
    "               batch_size=64, n_epochs=n_epochs, add_epochs=add_epochs)\n",
    "add_epochs += n_epochs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "## Plot Losses"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2017-10-24T12:29:58.654665",
     "start_time": "2017-10-24T12:29:57.159579Z"
    },
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def plot_losses(losses):\n",
    "    f = plt.figure()\n",
    "    #plt.plot(losses['dis_fake_loss'], label='dis_fake_loss')\n",
    "    #plt.plot(losses['dis_fake_loss'], label='dis_true_loss')\n",
    "    plt.plot(np.array(losses['gan'])[:,2], label='gan loss')\n",
    "    plt.legend()\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2017-10-24T12:30:00.128749",
     "start_time": "2017-10-24T12:29:58.660665Z"
    }
   },
   "outputs": [],
   "source": [
    "plot_losses(losses)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2017-10-24T12:29:05.837644",
     "start_time": "2017-10-24T12:29:04.428563Z"
    },
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "np.array(losses['gan'])[:,0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "# Generate gif"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2017-10-24T16:29:53.382088",
     "start_time": "2017-10-24T16:29:44.050554"
    }
   },
   "outputs": [],
   "source": [
    "import imageio\n",
    "import os\n",
    "RES_DIR = os.path.join(os.pardir, os.pardir, 'data', 'deep_learning')\n",
    "dir_path = os.path.join(RES_DIR, 'data', 'mnist_wgan', 'test_2')\n",
    "filenames = [(os.path.join(dir_path, filename), int(filename[9:-4])) for filename in os.listdir(dir_path)]\n",
    "images = []\n",
    "for filename in sorted(filenames, key=lambda x:x[1]):\n",
    "    images.append(imageio.imread(filename[0]))\n",
    "imageio.mimsave(os.path.join(RES_DIR, 'data', 'mnist_wgan', 'wgan.gif'), images)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "kernelspec": {
   "display_name": "Python [conda env:image-processing]",
   "language": "python",
   "name": "conda-env-image-processing-py"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.7"
  },
  "toc": {
   "nav_menu": {
    "height": "49px",
    "width": "252px"
   },
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "toc_cell": false,
   "toc_position": {
    "height": "817px",
    "left": "0px",
    "right": "1484px",
    "top": "74.2333px",
    "width": "212px"
   },
   "toc_section_display": "block",
   "toc_window_display": true
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
